import os
import re
import numpy as np
import cv2 as cv
from PyCmpltrtok.common import *
from PyCmpltrtok.common_np import *
import matplotlib.pyplot as plt
import PyCmpltrtok.data.cifar10.load_cifar10 as cifar10
from trt_my_helper_class import MyTrtHelper
import datetime


if '__main__' == __name__:

    def _main():
        trt_path = '/home/asuspei/PycharmProjects/AsusCondaP37Torch1101Cuda111/python_ai/category/onnx/torch2onnx/_save/vgg16_torch2onnx.py/v1.0/2022_06_17_12_54_40_476426-64.trt'
        # trt_path = '/home/asuspei/PycharmProjects/AsusCondaP37Torch1101Cuda111/python_ai/category/onnx/torch2onnx/_save/vgg16_torch2onnx.py/v1.0/2022_06_17_12_54_40_476426-64.fp16.trt'
        # trt_path = '_models/2022_06_17_12_54_40_476426-64.trt'
        # trt_path = '/mnt/d/_const/wsl/svn/python_ai/category/onnx/torch2onnx/_model/2022_06_17_12_54_40_476426-64.trt'
        TMP = 0
        PRECISION = np.float32

        sep('Load test dataset')
        x_train, y_train, x_test, y_test, label_names = cifar10.load()
        shape_ = cifar10.shape_
        print('label_names', label_names)
        print([(i, name) for i, name in enumerate(label_names)])
        if TMP:
            x_train = x_train[:512]
            y_train = y_train[:512]
            x_test = x_test[:512]
            y_test = y_test[:512]
        x_test = x_test.reshape(-1, *shape_)
        x_test = uint8_to_flt_by_lut(x_test, dtype=PRECISION)
        check_np(x_test, 'x_test')
        check_np(y_test, 'y_test')
        M = len(x_test)
        BATCH = 64
        N_BATCH = int(np.ceil(M / BATCH))
        print('N_BATCH', N_BATCH)

        sep('Load model by TRT onnx helper')
        print(trt_path)
        dt1 = datetime.datetime.now()
        myObj = MyTrtHelper(trt_path, PRECISION, [BATCH, *shape_], [BATCH, 10], output_dtype=PRECISION)
        dt2 = datetime.datetime.now()
        print('Init Duration', dt2 - dt1)

        sep('Run it')
        y_pred = None
        dt1 = datetime.datetime.now()
        for i in range(N_BATCH):
            bx = x_test[i*BATCH:(i+1)*BATCH]
            input_tpl = np.zeros([BATCH, *bx.shape[1:]], dtype=bx.dtype)
            input_tpl[:len(bx)] = bx
            # print(myObj.input_shape)  # tmp
            # print(input_tpl.shape)  # tmp
            output_arr = myObj.predict(input_tpl)
            if y_pred is None:
                y_pred = output_arr.copy()  # ATTENTION: .copy() is critical for ONNXClassifierWrapper.predict
            else:
                y_pred = np.concatenate([y_pred, output_arr], axis=0)
            print(i, end=', ', flush=True)
        print()
        dt2 = datetime.datetime.now()
        print('Infer all batches Duration', dt2 - dt1)
        check_np(y_pred, 'y_pred')
        y_pred = y_pred.argmax(axis=1)
        y_pred = y_pred[:M]
        check_np(y_pred, 'y_pred')
        check_np(y_test, 'y_test')
        acc = (y_pred == y_test).astype(np.uint8).sum() / M
        print('acc', acc)

        sep('Plot')
        plt.figure(figsize=(14, 6))
        plt.subplots_adjust(wspace=0.6)
        spr = 6
        spc = 12
        spn = 0
        m, r = spr * spc, 0
        base = 0
        print(f'Offet={base}')
        for j in range(m):
            spn += 1
            i = base + j
            if i >= M:
                break
            plt.subplot(spr, spc, spn)
            plt.axis('off')
            gt = y_test[i]
            h = y_pred[i]
            is_right = gt == h
            if is_right:
                r += 1
            h_label = label_names[h]
            plt.title(h_label if is_right else f'{h_label}(g:{label_names[gt]})', color='black' if is_right else 'red')
            img = x_test[i].transpose(1, 2, 0).astype(np.float32)
            plt.imshow(img)
        print(f'Accuracy of demo: {r/m:.4f}')
        print('Check the plotting window and close it to continue ...')
        plt.show()

    _main()
